package iot.queue;

import org.springframework.scheduling.concurrent.ThreadPoolTaskExecutor;
import org.springframework.util.Assert;

import java.util.ArrayList;
import java.util.List;
import java.util.concurrent.LinkedBlockingQueue;
import java.util.concurrent.atomic.AtomicInteger;
import java.util.function.Consumer;

/**
 * Created by dave on 18-7-26 下午11:37.
 * 入库数据队列管理类。建立了一个写入管理线程，负责管理从队列缓冲区中取出数据写入数据库。
 */
public class DataQueue<T> {
    private LinkedBlockingQueue<T> queue;
    private int poolSize;
    private int batchSize;
    private int concurrentWrite;
    private Class<T> entityType;
    private ThreadPoolTaskExecutor executor;
    private Consumer<List<T>> batchSaveFunc;
    private AtomicInteger queueSize = new AtomicInteger(0); // 当前队列中的数据条数
    private AtomicInteger writeThread = new AtomicInteger(0); // 当前正在进行的写入任务线程数

    public DataQueue(Class<T> entityType, Consumer<List<T>> batchSaveFunc, ThreadPoolTaskExecutor executor) {
        this(3000, 100, 5, entityType, batchSaveFunc, executor);
    }

    /**
     *
     * @param poolSize 缓冲池大小。比如3000。
     * @param batchSize 批量写入最大一次写入条数，比如100。
     * @param concurrentWrite 并发写数量，比如5。
     * @param entityType 实体类类型
     * @param executor 写入数据库任务线程池
     * @param batchSaveFunc 批量写入数据库方法。此方法在本DataQueue实例的内部线程中执行。
     */
    public DataQueue(int poolSize, int batchSize, int concurrentWrite, Class<T> entityType, Consumer<List<T>> batchSaveFunc, ThreadPoolTaskExecutor executor) {
        Assert.isTrue(poolSize > batchSize, "缓冲池大小必须大于批量写入大小");
        this.queue = new LinkedBlockingQueue<>(poolSize);
        this.poolSize = poolSize;
        this.batchSize = batchSize;
        this.concurrentWrite = concurrentWrite;
        this.entityType = entityType;
        this.batchSaveFunc = batchSaveFunc;
        this.executor = executor;
    }

    /**
     * 向数据队列添加数据。如果列表已满返回false。
     * @param data
     * @return
     */
    public boolean offer(T data) {
        if (queue.offer(data)) {
            int size = queueSize.incrementAndGet();
            if (size >= batchSize && writeThread.get() < concurrentWrite) {
                doSave();
            }
            return true;
        } else {
            if (doSave() > 0) { // 立即请求写入数据库后再试一次
                return offer(data);
            }
            return false;
        }
    }

    // 当前队列是否已满
    public boolean isFull() {
        return queueSize.get() >= poolSize;
    }

    // DataQueueDispatcher每100ms会调用此方法
    void onTimer() {
        doSave();
    }

    private int doSave() {
        if (writeThread.get() < concurrentWrite) {
            List<T> list = new ArrayList<>(batchSize);
            int size = queue.drainTo(list, batchSize);
            if (size > 0) {
                queueSize.addAndGet(-size);
                writeThread.incrementAndGet();
                executor.submitListenable(() -> batchSaveFunc.accept(list)).completable().thenApplyAsync(ret -> {
                    writeThread.decrementAndGet();
                    if (queueSize.get() >= batchSize) { // 继续写入
                        doSave();
                    }
                    return null;
                });
            }
            return size;
        } else {
            return 0;
        }
    }

    public int getPoolSize() {
        return poolSize;
    }

    public int getQueueSize() {
        return queueSize.get();
    }

    public Class<T> getEntityType() {
        return entityType;
    }

    public int getConcurrentWrite() {
        return concurrentWrite;
    }
}
